fmha_out = None
        # NOTE: (changwenbin) Bring out the public calculation in PD MIX to avoid repeated calculation.
        query = self.q_a_proj(hidden_states)
        query = self.q_a_layernorm(query)
        query = self.q_b_proj(query)
        query = query.reshape([-1, self.num_attention_heads_tp, self.qk_head_dim])
        query_nope, query_pe = query.split([self.qk_nope_head_dim, self.qk_rope_head_dim], axis=-1)
        compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
        compressed_kv, key_pe = compressed_kv.split([self.kv_lora_rank, self.qk_rope_head_dim], axis=-1)
        key_pe = key_pe.reshape([-1, 1, self.qk_rope_head_dim])
        compressed_kv = self.kv_a_layernorm(compressed_kv)
        query_pe, key_pe = self.rotary_emb(position_ids, query_pe, key_pe)